package org.example;

import com.sun.jna.Pointer;
import org.example.library.WintunLibrary;
import org.pcap4j.packet.*;
import org.pcap4j.packet.namednumber.IpNumber;
import org.pcap4j.packet.namednumber.IpVersion;
import org.pcap4j.packet.namednumber.TcpPort;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.*;
import java.util.Arrays;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

public class TcpHandler {

    private Pointer pointer;

    private ConcurrentHashMap<String, TcpSocket> hashMap = new ConcurrentHashMap<>();

    //虚拟线程
    private ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor();

    public TcpHandler(Pointer sessionHandle) {
        this.pointer = sessionHandle;
    }


    public void parsePacket(IpPacket packet) {
        try {
            TcpPacket tcpPacket = (TcpPacket) packet.getPayload();

            boolean syn = tcpPacket.getHeader().getSyn();
            boolean fin = tcpPacket.getHeader().getFin();
            boolean psh = tcpPacket.getHeader().getPsh();
            boolean ack = tcpPacket.getHeader().getAck();
            boolean rst = tcpPacket.getHeader().getRst();

            InetAddress dstAddr = packet.getHeader().getDstAddr();
            int dstPort = tcpPacket.getHeader().getDstPort().valueAsInt();

            int sequenceNumber = tcpPacket.getHeader().getSequenceNumber();


            InetAddress srcAddr = packet.getHeader().getSrcAddr();
            int srcPort = tcpPacket.getHeader().getSrcPort().valueAsInt();

            String key = srcAddr + "#" + srcPort;
            TcpSocket tcpSocket = hashMap.get(key);
            //第一次握手
            if (syn) {
                if (tcpSocket == null) {
                    tcpSocket = new TcpSocket(dstAddr, dstPort,packet,sequenceNumber);
                    hashMap.put(key, tcpSocket);
                    tcpSocket.connection();
                }
            }else {

                //服务重启导致的
                if(tcpSocket == null){
                    new TcpSocket(dstAddr, dstPort,packet,sequenceNumber).replyRst(tcpPacket);
                    return;
                }

                //连接重置
                if(rst){
                    hashMap.remove(key);
                    tcpSocket.close();
                    return;
                }

                //第四次挥手
                if(ack && tcpSocket.sendFin.get() && tcpSocket.replyFin.get()){
                    hashMap.remove(key);
                    tcpSocket.close();
                    return;
                }
                //写入队列
                tcpSocket.writeQueue.offer(tcpPacket);

            }
        }catch (Exception ex){
            ex.printStackTrace();
        }

    }

    class TcpSocket {

        private  AtomicInteger seq = new AtomicInteger(1000);
        private AtomicInteger ack;

        private InetAddress dstAddr;
        private int dstPort;
        private Socket socket;
        private IpPacket srcIpPacket;
        private AtomicBoolean sendFin = new AtomicBoolean();
        private AtomicBoolean replyFin = new AtomicBoolean();


        private BlockingQueue<TcpPacket> writeQueue = new LinkedBlockingQueue<>();


        public TcpSocket(InetAddress dstAddr, int dstPort, IpPacket ipPacket, int ack) {
            this.dstAddr = dstAddr;
            this.dstPort = dstPort;
            this.srcIpPacket = ipPacket;
            this.ack = new AtomicInteger(ack+1);
        }

        public void connection() {
            try {
                Proxy proxy = new Proxy(Proxy.Type.SOCKS, new InetSocketAddress("127.0.0.1", 9000));
                socket = new Socket(proxy);
//                socket.setSoTimeout(5000);
                //转换域名
                String domain = SimpleDnsServer.biMap.getKey(dstAddr);
                //不需要解析域名
                SocketAddress remoteAddr = InetSocketAddress.createUnresolved(domain, dstPort);
                socket.connect(remoteAddr);

                replySyn();

                //读任务
                readTask();
                //写任务
                writeTask();
            } catch (IOException e) {
                e.printStackTrace();
                try {
                    socket.close();
                } catch (IOException ex) {
                    throw new RuntimeException(ex);
                }
            }
        }

        private void readTask() {
            executor.submit(()->{
                try (InputStream inputStream = socket.getInputStream()) {
                    byte[] buffer = new byte[1024*1024];
                    int length;
                    while ((length = inputStream.read(buffer)) != -1) {
                        // 处理接收到的数据
                        byte[] bytes = Arrays.copyOfRange(buffer, 0, length);
                        replyPsh(bytes);
                    }
                    System.out.println("客户端关闭连接");
                } catch (IOException e) {
//                    e.printStackTrace();
                }finally {
                    if(!sendFin.getAndSet(true)){
                        replyFin();
                    }
                }
            });
        }

        private void writeTask() {
            executor.submit(()->{
                try (OutputStream outputStream = socket.getOutputStream()){
                    while (true) {
                        TcpPacket tcpPacket = writeQueue.take();
                        Packet payload = tcpPacket.getPayload();
                        //有数据就写入socket
                        if (payload != null && payload.getRawData().length > 0) {
                            byte[] rawData = payload.getRawData();
                            outputStream.write(rawData);
                            ack.addAndGet(rawData.length);
                            replyAck();
                        }
                        if(tcpPacket.getHeader().getFin()){
                            this.ack.incrementAndGet();
                            this.replyAck();
                            this.replyFin.set(true);

                            if(!this.sendFin.getAndSet(true)){
                                this.replyFin();
                            }
                        }
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                }
            });
        }

        private void close(){
            if(socket!=null){
                try {
                    socket.close();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
        }

        /**
         * 写入网络适配器
         * @param data
         */
        private void writeData(byte[] data) {

            Pointer packetPointer = WintunLibrary.INSTANCE.WintunAllocateSendPacket(pointer, data.length);

            if (packetPointer != null) {

                packetPointer.write(0, data, 0, data.length);

                // 发送数据包
                WintunLibrary.INSTANCE.WintunSendPacket(pointer, packetPointer);

            } else {
                System.err.println("Failed to allocate send packet.");
            }
        }

        /** 下面构建TCP数据包   **/

        private void replyRst(TcpPacket tcpPacket) {
            int sequenceNumber = tcpPacket.getHeader().getSequenceNumber();
            int acknowledgmentNumber = tcpPacket.getHeader().getAcknowledgmentNumber();
            seq.set(acknowledgmentNumber);
            ack.set(sequenceNumber);
            buildIpPack(null,16);
        }

        public void replySyn() {

            IpPacket newIpacket = buildIpPack(null,1 | 8);
            byte[] rawData = newIpacket.getRawData();
            writeData(rawData);

            this.seq.incrementAndGet();
        }

        public void replyFin() {


            IpPacket newIpacket = buildIpPack(null,4 | 8);
            byte[] rawData = newIpacket.getRawData();
            writeData(rawData);

            this.seq.incrementAndGet();
        }

        public void replyAck() {

            IpPacket newIpacket = buildIpPack(null,8);
            byte[] rawData = newIpacket.getRawData();

            writeData(rawData);
        }

        public void replyPsh(byte[] data) {

            IpPacket newIpacket = buildIpPack(data,2 | 8);
            byte[] rawData = newIpacket.getRawData();

            writeData(rawData);

            seq.getAndAdd(data.length);

        }
        public IpPacket buildIpPack(byte[] data, int flags) {

            TcpPacket localtcpPacket = (TcpPacket) srcIpPacket.getPayload();
            TcpPort srcPort = localtcpPacket.getHeader().getSrcPort();
            TcpPort dstPort = localtcpPacket.getHeader().getDstPort();


            InetAddress srcAddr = srcIpPacket.getHeader().getSrcAddr();
            InetAddress dstAddr = srcIpPacket.getHeader().getDstAddr();

            TcpPacket.Builder tcpBuilder = new TcpPacket.Builder()
                    .srcPort(dstPort)
                    .dstPort(srcPort)
                    .sequenceNumber(this.seq.get())
                    .acknowledgmentNumber(this.ack.get())
                    .window(Short.MIN_VALUE)
                    .syn((flags & 1) == 1)
                    .psh((flags & 2) == 2)
                    .fin((flags & 4) == 4)
                    .ack((flags & 8) == 8)
                    .rst((flags & 16) == 16)
                    .srcAddr(dstAddr)
                    .dstAddr(srcAddr)
                    .correctChecksumAtBuild(true)
                    .correctLengthAtBuild(true);

            if (data != null && data.length > 0) {
                UnknownPacket.Builder unknownBuilder = new UnknownPacket.Builder();
                unknownBuilder.rawData(data);
                tcpBuilder.payloadBuilder(unknownBuilder);
            }

            if (srcAddr instanceof Inet6Address) {
                return buildIpv6Packet(tcpBuilder, srcIpPacket);
            } else {
                return buildIpv4Packet(tcpBuilder, srcIpPacket);
            }
        }

        private  IpPacket buildIpv6Packet(TcpPacket.Builder tcpBuilder, IpPacket srcIpPacket) {
            InetAddress srcAddr = srcIpPacket.getHeader().getSrcAddr();
            InetAddress dstAddr = srcIpPacket.getHeader().getDstAddr();

            return new IpV6Packet.Builder()
                    .version(IpVersion.IPV6)
                    .srcAddr((Inet6Address)dstAddr)
                    .dstAddr((Inet6Address)srcAddr)
                    .trafficClass(IpV6SimpleTrafficClass.newInstance((byte) 0))
                    .flowLabel(IpV6SimpleFlowLabel.newInstance(0))
                    .nextHeader(IpNumber.TCP)
                    .correctLengthAtBuild(true)
                    .payloadBuilder(tcpBuilder)
                    .build();
        }

        private  IpPacket buildIpv4Packet(TcpPacket.Builder tcpBuilder, IpPacket srcIpPacket) {
            InetAddress srcAddr = srcIpPacket.getHeader().getSrcAddr();
            InetAddress dstAddr = srcIpPacket.getHeader().getDstAddr();
            return new IpV4Packet.Builder()
                    .version(IpVersion.IPV4)
                    .ttl((byte) 45)
                    .tos(IpV4Rfc1349Tos.newInstance((byte) 0))
                    .protocol(IpNumber.TCP)
                    .srcAddr((Inet4Address) dstAddr)
                    .dstAddr((Inet4Address) srcAddr)
                    .correctLengthAtBuild(true)
                    .correctChecksumAtBuild(true)
                    .payloadBuilder(tcpBuilder)
                    .identification((short) ipId.incrementAndGet())
                    .build();
        }

        private static AtomicInteger ipId = new AtomicInteger(0);

    }

}
